package org.terracotta.workmanager.routing;

import com.google.zahra.testmanager.Test;

import java.util.Map;
import java.util.Set;

import org.terracotta.workmanager.pipemanager.Pipe;
import org.terracotta.workmanager.pipemanager.PipeManager;

import commonj.work.Work;
import commonj.work.WorkEvent;
import commonj.work.WorkException;
import commonj.work.WorkListener;

/**
 * Load balancing router, routes work to the shortest queue.
 */
public class LoadBalancingRouter<ID> implements Router<ID> {

	private final PipeManager<ID> pipeManager;

	private final RoutableWorkItem.Factory<ID> m_workItemFactory = new RoutableWorkItem.Factory<ID>();

	public LoadBalancingRouter(final PipeManager<ID> pipeManager) {
		this.pipeManager = pipeManager;
	}

	public RoutableWorkItem<ID> route(final Work work) {
		return route(work, null);
	}

	public RoutableWorkItem<ID> route(final Work work, final WorkListener listener) {
		waitForAtLeastOnePipe();
		
		LoadBalancingRouter.PipeLength<ID> shortestQueue;
		
		if (((Test)work).os != 0) {
		  shortestQueue = getShortestWorkQueueForOs(((Test)work).os);
		}
		else {
		  shortestQueue = getShortestWorkQueue();
	    }
		
		RoutableWorkItem<ID> workItem = m_workItemFactory.create(work, listener, shortestQueue.routingID);
        try {
          System.out.println("Routing " + ((Test)workItem.getResult()).jobID + " to " + shortestQueue.routingID );
            shortestQueue.pipe.getPendingWorkQueue().put(workItem);
            workItem.setStatus(WorkEvent.WORK_SCHEDULED, null);
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
        }
        isInterrupted(workItem);
        return workItem;
	}

	public RoutableWorkItem<ID> route(final RoutableWorkItem<ID> workItem) {
		waitForAtLeastOnePipe();		
		
		LoadBalancingRouter.PipeLength<ID> shortestQueue;
		Test test = (Test)workItem.getResult();
        
        if (test.os != 0) {
          shortestQueue = getShortestWorkQueueForOs(test.os);
        }
        else {
          shortestQueue = getShortestWorkQueue();
        }
		
		synchronized (workItem) {
			workItem.setRoutingID(shortestQueue.routingID);
		}
		try {
		  System.out.println("Routing " + ((Test)workItem.getResult()).jobID + " to " + shortestQueue.routingID );
			shortestQueue.pipe.getPendingWorkQueue().put(workItem);
			workItem.setStatus(WorkEvent.WORK_SCHEDULED, null);
		} catch (InterruptedException e) {
			Thread.currentThread().interrupt();
		}
		isInterrupted(workItem);
		return workItem;
	}

    public Pipe<RoutableWorkItem<ID>, ID> getPipeFor(ID routingID) {
      return pipeManager.getOrCreatePipeFor(routingID);
    }
    
    public Pipe<RoutableWorkItem<ID>, ID> getPipeForOrNull(ID routingID) {
      return pipeManager.get(routingID);
    }
  
	private void isInterrupted(RoutableWorkItem<ID> workItem) {
		if (Thread.interrupted()) {
			workItem.setStatus(WorkEvent.WORK_REJECTED, new WorkException("thread interrupted"));
		}
	}

	private synchronized void waitForAtLeastOnePipe() {
		try {
			while (pipeManager.getPipes().isEmpty()) {
				System.out.println("waiting for at least one worker to join");
				synchronized (pipeManager.getPipes()) {
				  pipeManager.getPipes().wait();
				}
			}
		} catch (InterruptedException e) {
			Thread.currentThread().interrupt();
		}
	}
	
    private LoadBalancingRouter.PipeLength<ID> getShortestWorkQueueForOs(int os) {
      LoadBalancingRouter.PipeLength<ID> shortestQueue = new LoadBalancingRouter.PipeLength<ID>();

      int queueLengthForShortestQueue = Integer.MAX_VALUE;
      ID routingIDForShortestQueue = null;
      Map<ID, Pipe<RoutableWorkItem<ID>, ID>> allPipes = pipeManager.getPipesForOs(os);
      
      for (Map.Entry<ID, Pipe<RoutableWorkItem<ID>, ID>> pipes : allPipes.entrySet()) {
          ID routingID = pipes.getKey();
          Pipe<RoutableWorkItem<ID>, ID> pipe = pipes.getValue();
          int queueSize = pipe.getPendingWorkQueue().size();
          if (queueSize <= queueLengthForShortestQueue) {
              queueLengthForShortestQueue = queueSize;
              routingIDForShortestQueue = routingID;
          }
      }
      shortestQueue.pipe = pipeManager.getOrCreatePipeFor(routingIDForShortestQueue);
      shortestQueue.routingID = routingIDForShortestQueue;
      return shortestQueue;
  }

	private LoadBalancingRouter.PipeLength<ID> getShortestWorkQueue() {
		LoadBalancingRouter.PipeLength<ID> shortestQueue = new LoadBalancingRouter.PipeLength<ID>();

		int queueLengthForShortestQueue = Integer.MAX_VALUE;
		ID routingIDForShortestQueue = null;
		Map<ID, Pipe<RoutableWorkItem<ID>, ID>> allPipes = pipeManager.getPipes();
		for (Map.Entry<ID, Pipe<RoutableWorkItem<ID>, ID>> pipes : allPipes.entrySet()) {
			ID routingID = pipes.getKey();
			Pipe<RoutableWorkItem<ID>, ID> pipe = pipes.getValue();
			int queueSize = pipe.getPendingWorkQueue().size();
			if (queueSize <= queueLengthForShortestQueue) {
				queueLengthForShortestQueue = queueSize;
				routingIDForShortestQueue = routingID;
			}
		}
		shortestQueue.pipe = pipeManager.getOrCreatePipeFor(routingIDForShortestQueue);
		shortestQueue.routingID = routingIDForShortestQueue;
		return shortestQueue;
	}

	public void reset() {
		pipeManager.reset();
	}

	public Set<RoutableWorkItem<ID>> getAllCompletedWork() {
		return pipeManager.getAllCompletedWork();
	}

	public synchronized void register(ID routingID) {
		pipeManager.getOrCreatePipeFor(routingID);
		notifyAll();
	}

	public void unregister(ID routingID) {
	  System.out.println("Unregistering the pipe with ID: " + routingID);
		pipeManager.removePipeFor(routingID);
	}

	private static class PipeLength<ID> {
		public ID routingID;

		public Pipe<RoutableWorkItem<ID>, ID> pipe;

		public int pipeLength = 1000;
	}
}
